/*
 * (c) 2008-2011 Daniel Halperin <dhalperi@cs.washington.edu>
 */
#include "iwl_connector.h"
// #include "common.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <linux/netlink.h>

#define MAX_PAYLOAD 2048
#define SLOW_MSG_CNT 1

int socket_fd = -1;                  // The socket file descriptor
FILE* output_file = NULL;

void exit_program(int code)
{
    if (output_file)
    {
        fclose(output_file);
        output_file = NULL;
    }
    if (socket_fd != -1)
    {
        close(socket_fd);
        socket_fd = -1;
    }
    exit(code);
}

void exit_program_with_error(int code, char* msg)
{
    perror(msg);
    exit_program(code);
}

void check_usage(int argc, char** argv)
{
    if (argc != 2)
    {
        fprintf(stderr, "Usage: %s <output_file>\n", argv[0]);
        exit_program(1);
    }
}

FILE* open_output_file(char* filename, char* mode)
{
    FILE* file_ptr = fopen(filename, mode);
    if (!file_ptr)
    {
        perror("fopen");
        exit_program(1);
    }
    return file_ptr;
}

void handle_signal(int sig)
{
    fprintf(stderr, "Caught signal %d\n", sig);
    exit_program(0);
}


int main(int argc, char** argv)
{
    /* Local variables */
    // struct sockaddr_nl process_addr, kernel_addr; // Addresses for recv, send, bind
    struct sockaddr_nl process_addr; // Addresses for recv, send, bind
    struct cn_msg *cn_msg_ptr;
    char buffer[4096];
    int result;
    unsigned short payload_len, payload_len_nw;
    int msg_count = 0;

    /* Ensure correct usage */
    check_usage(argc, argv);

    /* Open and check output file */
    output_file = open_output_file(argv[1], "w");

    /* Set up the socket */
    socket_fd = socket(PF_NETLINK, SOCK_DGRAM, NETLINK_CONNECTOR);
    if (socket_fd == -1)
        exit_program_with_error(-1, "socket");

    /* Initialize the address structs */
    memset(&process_addr, 0, sizeof(struct sockaddr_nl));
    process_addr.nl_family = AF_NETLINK;
    process_addr.nl_pid = getpid(); // This process' PID
    process_addr.nl_groups = CN_IDX_IWLAGN;
    // memset(&kernel_addr, 0, sizeof(struct sockaddr_nl));
    // kernel_addr.nl_family = AF_NETLINK;
    // kernel_addr.nl_pid = 0; // Kernel
    // kernel_addr.nl_groups = CN_IDX_IWLAGN;

    /* Bind the socket */
    if (bind(socket_fd, (struct sockaddr *)&process_addr, sizeof(struct sockaddr_nl)) == -1)
        exit_program_with_error(-1, "bind");

    /* Subscribe to the netlink group */
    {
        int on = process_addr.nl_groups;
        result = setsockopt(socket_fd, 270, NETLINK_ADD_MEMBERSHIP, &on, sizeof(on));
        if (result)
            exit_program_with_error(-1, "setsockopt");
    }

    /* Set up the "handle_signal" function as this program's signal handler */
    signal(SIGINT, handle_signal);

    /* Continuously poll the socket waiting for a message */
    while (1)
    {
        /* Receive from the socket with infinite timeout */
        result = recv(socket_fd, buffer, sizeof(buffer), 0);
        if (result == -1)
            exit_program_with_error(-1, "recv");

        /* Extract the message portion and print some stats */
        cn_msg_ptr = NLMSG_DATA(buffer);
        if (msg_count % SLOW_MSG_CNT == 0)
            printf("received %d bytes: id: %d val: %d seq: %d clen: %d\n", 
            cn_msg_ptr->len, cn_msg_ptr->id.idx, cn_msg_ptr->id.val, cn_msg_ptr->seq, cn_msg_ptr->len);

        /* Log the data to the output file */
        payload_len = (unsigned short) cn_msg_ptr->len;
        payload_len_nw = htons(payload_len);
        fwrite(&payload_len_nw, 1, sizeof(unsigned short), output_file);
        result = fwrite(cn_msg_ptr->data, 1, payload_len, output_file);

        if (msg_count % 100 == 0)
            printf("wrote %d bytes [msgcnt=%u]\n", result, msg_count);
    
        ++msg_count;

        if (result != payload_len)
            exit_program_with_error(1, "fwrite");
    }

    exit_program(0);
    return 0;
}
